# model of EfficientNet B0 ~ B7
# Ddddavid 2021/08/22
# Python 3.8  Tensorflow 2.3.0


import os
import numpy as np
import math
import tensorflow as tf
from efficientnet import EfficientNet, efficient_net_b7, get_efficient_net


BATCH_SIZE = 16
EPOCHS = 10
model = efficient_net_b7()


images, labels = [], []


x_tr, y_tr = None, None
x_te, y_te = None, None


loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.layers.RMSprop()

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalCrossentropy(name='train_accuracy')

valid_loss = tf.keras.metrics.Mean(name='valid_loss')
valid_accuracy = tf.keras.metrics.SparseCategoricalCrossentropy(name='valid_accuracy')



def loss(predict, real):
    return tf.reduce_mean(tf.pow(predict - real, 2))


def train_step(image_batch, label_batch):
    with tf.GradientTape() as tape:
        prefictions = model(image_batch, training=True)
        loss = loss_object(y_true=label_batch, y_pred=predictions)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(grads_and_vars=zip(gradients, model.trainable_variables))
    
    train_loss.update_state(values=loss)
    train_accuracy.update_state(y_true=label_batch, y_pred=predictions)
    

def valid_step(image_batch, label_batch):
    predictions = model(image_batch, training=False)
    valid_loss = loss_object(label_batch, predictions)
    
    valid_loss.update_state(values=valid_loss)
    valid_accuracy.update_state(y_true=label_batch, y_pred=predictions)
    
    
    

for epoch in range(EPOCHS):
    step = 0
    for features in train_dataset:
        train_images, train_labels = process_features(features, data_augmentation=True)
        train_step(train_images, train_labels)
        
        print('Epoch {}/{}, step = {}/{},'
              'loss = {:.5f}, acc = {:.5f}'.format(
                  epoch, EPOCHS, step,
                  math.ceil(train_count / BATCH_SIZE),
                  train_loss.result().numpy(),
                  train_accuracy.result().numpy()
              ))

    for features in valid_dataset:
        valid_images, valid_labels = process_features(features, data_augmentation=False)
        valid_step(valid_image, valid_labels)
        
    print("Epoch {}/{}, train loss = {:.5f}, train acc = {:.5f}," 
                "valid loss {:.5}, valid acc = {:.5}".format(
                    epoch, EPOCHS, train_loss.result().numpy(),
                    valid_accuracy.result().numpy(),
                    valid_loss.result().numpy(),
                    valid_accuracy.result().numpy()
                ))









